/**
 * @file test_rstream.cpp
 *
 * Test random streams
 * 
 * @author Dahua Lin
 */


#include <gtest/gtest.h>

#include <slipp/base/math_base.h>
#include <slipp/prob_base/rstream.h>

#include "eigen_test.h"
#include "random_test.h"

using namespace sli;
using namespace sli::test;

typedef std_rstream<default_std_rand_eng> rstream_t;

TEST( StdStream, RandI32 )
{
	default_std_rand_eng eng;

	rstream_t rs1(eng);
	rstream_t rs2(eng);

	int32_t m = 10;
	const index_t N = 1000001;

	iarow_t s1(N);
	for (index_t i = 0; i < N; ++i)
	{
		s1(i) = rs1.rand_i32(m);
	}

	iarow_t s2(N);
	s2.fill(0);
	rs2.rand_i32_m(m, (size_t)N, s2.data());

	ASSERT_TRUE( test_equal(s1, s2) );
	EXPECT_EQ( s1.minCoeff(), 0 );
	EXPECT_EQ( s1.maxCoeff(), m-1 );

	darow_t p = darow_t::Constant(m, 1.0 / m);
	EXPECT_TRUE( histogram_check(s1, p, 0, 2.0e-3) );
}


TEST( StdStream, RandU32 )
{
	default_std_rand_eng eng;

	rstream_t rs1(eng);
	rstream_t rs2(eng);

	uint32_t m = 10;
	const index_t N = 1000001;

	uarow_t s1(N);
	for (index_t i = 0; i < N; ++i)
	{
		s1(i) = rs1.rand_u32(m);
	}

	uarow_t s2(N);
	s2.fill(0);
	rs2.rand_u32_m(m, (size_t)N, s2.data());

	ASSERT_TRUE( test_equal(s1, s2) );
	EXPECT_EQ( s1.minCoeff(), 0 );
	EXPECT_EQ( s1.maxCoeff(), m-1 );

	darow_t p = darow_t::Constant(m, 1.0 / m);
	EXPECT_TRUE( histogram_check(s1, p, 0, 2.0e-3) );
}


TEST( StdStream, RandF64 )
{
	default_std_rand_eng eng;

	rstream_t rs1(eng);
	rstream_t rs2(eng);

	const index_t N = 1000001;

	darow_t s1(N);
	for (index_t i = 0; i < N; ++i)
	{
		s1(i) = rs1.rand_f64();
	}

	darow_t s2(N);
	s2.fill(0);
	rs2.rand_f64_m((size_t)N, s2.data());

	ASSERT_TRUE( test_equal(s1, s2) );
	EXPECT_TRUE( (s1 > 0.0).all() );
	EXPECT_TRUE( (s1 < 1.0).all() );
	EXPECT_NEAR( s1.mean(), 0.5, 1.0e-3 );
	EXPECT_NEAR( s1.minCoeff(), 0.0, 1.0e-3 );
	EXPECT_NEAR( s1.maxCoeff(), 1.0, 1.0e-3 );
	EXPECT_NEAR( (s1 - 0.5).square().mean(), 1.0/12, 5.0e-3 );
	EXPECT_NEAR( average_neighbor_diff(s1), 1.0/3, 1.0e-2 );
}


inline double nrmd_abs_cdf(double x)
{
	double cp = 0.5 * (1 + math::erf(SLI_M_SQRT1_2 * x));
	double cn = 0.5 * (1 + math::erf(SLI_M_SQRT1_2 * (-x)));
	return cp - cn;
}

inline double exp_cdf(double x)
{
	return 1 - math::exp(-x);
}


TEST( RNGs, BernoulliRng )
{
	default_std_rand_eng eng;

	rstream_t rs1(eng);
	rstream_t rs2(eng);

	bernoulli_rng<rstream_t> rng1(rs1, 0.4);
	bernoulli_rng<rstream_t> rng2(rs2, 0.4);

	const index_t N = 1000001;

	barow_t s1(N);
	for (index_t i = 0; i < N; ++i)
	{
		s1(i) = rng1.get();
	}

	barow_t s2(N);
	rng2.get_m((size_t)N, s2.data());

	ASSERT_TRUE( test_equal(s1, s2) );
	EXPECT_NEAR( double(s1.count()) / double(N),  0.4, 1e-4 );
}

TEST( RNGs, DiscreteRng )
{
	default_std_rand_eng eng;

	rstream_t rs1(eng);
	rstream_t rs2(eng);

	const size_t K = 3;
	const double p[K] = {0.3, 0.5, 0.2};

	discrete_rng<rstream_t> rng1(rs1, (size_t)K, p, 1.0);
	discrete_rng<rstream_t> rng2(rs2, (size_t)K, p, 1.0);

	index_t N = 1000000;

	iarow_t s1(N);
	for (index_t i = 0; i < N; ++i)
	{
		s1(i) = rng1.get();
	}

	iarow_t s2(N);
	rng2.get_m((size_t)N, s2.data());

	ASSERT_TRUE( test_equal(s1, s2) );
	EXPECT_EQ( s1.minCoeff(), 0 );
	EXPECT_EQ( s1.maxCoeff(), K-1 );
	EXPECT_NEAR( double((s1 == 0).count()) / double(N), p[0], 2.0e-3 );
	EXPECT_NEAR( double((s1 == 1).count()) / double(N), p[1], 2.0e-3 );
	EXPECT_NEAR( double((s1 == 2).count()) / double(N), p[2], 2.0e-3 );
}


TEST( RNGs, DiscreteRngEx )
{
	default_std_rand_eng eng;

	rstream_t rs1(eng);
	rstream_t rs2(eng);

	const size_t K = 3;
	const double w[K] = {0.6, 0.8, 0.2};

	discrete_rng<rstream_t> rng1(rs1, (size_t)K, w, 2.0);
	discrete_rng<rstream_t> rng2(rs2, (size_t)K, w, 2.0);

	index_t N = 1000000;

	iarow_t s1(N);
	for (index_t i = 0; i < N; ++i)
	{
		s1(i) = rng1.get();
	}

	iarow_t s2(N);
	rng2.get_m((size_t)N, s2.data());

	ASSERT_TRUE( test_equal(s1, s2) );
	EXPECT_EQ( s1.minCoeff(), 0 );
	EXPECT_EQ( s1.maxCoeff(), K );
	EXPECT_NEAR( double((s1 == 0).count()) / double(N), 0.3, 2.0e-3 );
	EXPECT_NEAR( double((s1 == 1).count()) / double(N), 0.4, 2.0e-3 );
	EXPECT_NEAR( double((s1 == 2).count()) / double(N), 0.1, 2.0e-3 );
	EXPECT_NEAR( double((s1 == 3).count()) / double(N), 0.2, 2.0e-3 );
}



TEST( RNGs, NormalRng)
{
	default_std_rand_eng eng;

	rstream_t rs1(eng);
	rstream_t rs2(eng);

	normal_rng<rstream_t> rng1(rs1);
	normal_rng<rstream_t> rng2(rs2);

	const index_t N = 1000001;

	darow_t s1(N);
	for (index_t i = 0; i < N; ++i)
	{
		s1(i) = rng1.get();
	}

	darow_t s2(N);
	s2.fill(0);
	rng2.get_m((size_t)N, s2.data());

	ASSERT_TRUE( test_equal(s1, s2) );
	EXPECT_NEAR( s1.mean(), 0.0, 5.0e-3 );
	EXPECT_NEAR( s1.square().mean(), 1.0, 5.0e-3 );
	EXPECT_NEAR( average_neighbor_diff(s1), SLI_M_2_SQRTPI, 5.0e-3 );
	EXPECT_NEAR( double((s1.abs() < 0.5).count()) / double(N), nrmd_abs_cdf(0.5), 8.0e-3);
	EXPECT_NEAR( double((s1.abs() < 1.0).count()) / double(N), nrmd_abs_cdf(1.0), 8.0e-3);
	EXPECT_NEAR( double((s1.abs() < 2.0).count()) / double(N), nrmd_abs_cdf(2.0), 8.0e-3);
}


TEST( RNGs, ExponentialRng)
{
	default_std_rand_eng eng;

	rstream_t rs1(eng);
	rstream_t rs2(eng);

	exponential_rng<rstream_t> rng1(rs1);
	exponential_rng<rstream_t> rng2(rs2);

	const index_t N = 1000001;

	darow_t s1(N);
	for (index_t i = 0; i < N; ++i)
	{
		s1(i) = rng1.get();
	}

	darow_t s2(N);
	s2.fill(0);
	rng2.get_m((size_t)N, s2.data());

	ASSERT_TRUE( (s1 > 0.0).all() );
	ASSERT_TRUE( test_equal(s1, s2) );
	EXPECT_NEAR( s1.mean(), 1.0, 5.0e-3 );
	EXPECT_NEAR( (s1 - 1.0).square().mean(), 1.0, 1.0e-2 );
	EXPECT_NEAR( average_neighbor_diff(s1), 1.0, 1.0e-2 );
	EXPECT_NEAR( double((s1 < 0.5).count()) / double(N), exp_cdf(0.5), 8.0e-3);
	EXPECT_NEAR( double((s1 < 1.0).count()) / double(N), exp_cdf(1.0), 8.0e-3);
	EXPECT_NEAR( double((s1 < 2.0).count()) / double(N), exp_cdf(2.0), 8.0e-3);
}







